package org.yinweichen.llm_dataset_backend.controller;

import org.yinweichen.llm_dataset_backend.DTO.BatchGenerateRequest;
import org.yinweichen.llm_dataset_backend.entity.LlmAnswer;
import org.yinweichen.llm_dataset_backend.entity.StandardQuestion;
import org.yinweichen.llm_dataset_backend.entity.StandardQuestionId;
import org.yinweichen.llm_dataset_backend.repository.LlmAnswerRepository;
import org.yinweichen.llm_dataset_backend.repository.StandardQuestionRepository;
import org.yinweichen.llm_dataset_backend.util.ModelApiClient;
import org.springframework.web.bind.annotation.*;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

@RestController
@RequestMapping("/api/llm")
public class LlmController {
    private final LlmAnswerRepository llmAnswerRepository;
    private final StandardQuestionRepository standardQuestionRepository;

    public LlmController(
                       LlmAnswerRepository llmAnswerRepository,
                       StandardQuestionRepository standardQuestionRepository) {
        this.llmAnswerRepository = llmAnswerRepository;
        this.standardQuestionRepository = standardQuestionRepository;
    }

    /**
     * 获取AI回答
     * @param questionId 标准问题ID
     */
    @PostMapping("/answer")//TODO:设定了版本
    public LlmAnswer getModelAnswer(@RequestParam Long questionId,
                                    @RequestParam String version) {
        // 获取标准问题内容
        StandardQuestion question = standardQuestionRepository.findById(new StandardQuestionId(questionId,version))
            .orElseThrow(() -> new RuntimeException("Question not found"));

        // 调用ModelApiClient获取AI回答
        String prompt="You are a helpful assistant in Computer Science, please answer the following question within 600 words.";
        String aiAnswer = ModelApiClient.getAnswer(prompt,question.getQuestion());

        // 创建并保存回答记录
        LlmAnswer answer = new LlmAnswer();
        answer.setModel(ModelApiClient.getCurrModel());
        answer.setQuestion(question);
        answer.setAnswer(aiAnswer);

        return llmAnswerRepository.save(answer);
    }

    /**
     * 批量获取AI回答
     * @param request 包含模型列表和问题ID列表的请求体
     * @return 生成的所有回答列表
     */
    @PostMapping("/generate")//TODO:设定了版本
    public List<LlmAnswer> batchGetModelAnswers(@RequestBody BatchGenerateRequest request) {
        // 参数校验
        if (request == null) {
            throw new IllegalArgumentException("Request cannot be null");
        }
        if (request.questionIds() == null || request.questionIds().isEmpty()) {
            throw new IllegalArgumentException("Question IDs cannot be null or empty");
        }
        if (request.modelNames() == null || request.modelNames().isEmpty()) {
            throw new IllegalArgumentException("Model names cannot be null or empty");
        }

        // 获取所有标准问题
        List<StandardQuestion> questions = new ArrayList<>();
        String version= request.version();
        for(Long questionId : request.questionIds()) {
            questions.add(standardQuestionRepository.findById(new StandardQuestionId(questionId,version))
                .orElseThrow(() -> new RuntimeException("Question not found with id: " + questionId)));
        }
        String prompt="You are a helpful assistant in Computer Science, please answer the following question within 600 words.";
        // 并行处理每个模型和问题的组合
        return request.modelNames().parallelStream()
            .flatMap(modelName -> questions.stream()
                .map(question -> {
                    // 调用ModelApiClient获取AI回答
                    ModelApiClient.setCurrModel(modelName);
                    String aiAnswer = ModelApiClient.getAnswer(prompt,question.getQuestion());

                    // 创建回答记录
                    LlmAnswer answer = new LlmAnswer();
                    answer.setModel(modelName);
                    answer.setQuestion(question);
                    answer.setAnswer(aiAnswer);

                    return answer;
                }))
            .map(llmAnswerRepository::save)
            .collect(Collectors.toList());
    }
}
